from channels.middleware import BaseMiddleware
from channels.db import database_sync_to_async
from urllib.parse import parse_qs
import jwt
from django.conf import settings


class JWTAuthMiddleware(BaseMiddleware):
    """
    JWT认证中间件，用于WebSocket连接
    """
    
    def __init__(self, inner):
        super().__init__(inner)
    
    async def __call__(self, scope, receive, send):
        # 只处理WebSocket连接
        if scope['type'] == 'websocket':
            # 从查询参数中获取token
            query_string = scope.get('query_string', b'').decode()
            query_params = parse_qs(query_string)
            token = query_params.get('token', [None])[0]
            
            if token:
                try:
                    # 验证JWT token
                    from rest_framework_simplejwt.tokens import UntypedToken
                    UntypedToken(token)
                    decoded_data = jwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"])
                    user = await self.get_user(decoded_data['user_id'])
                    scope['user'] = user
                except Exception:
                    from django.contrib.auth.models import AnonymousUser
                    scope['user'] = AnonymousUser()
            else:
                from django.contrib.auth.models import AnonymousUser
                scope['user'] = AnonymousUser()
        
        return await super().__call__(scope, receive, send)
    
    @database_sync_to_async
    def get_user(self, user_id):
        try:
            from django.contrib.auth import get_user_model
            User = get_user_model()
            return User.objects.get(id=user_id)
        except:
            from django.contrib.auth.models import AnonymousUser
            return AnonymousUser()
